package main

import (
	"flag"
	"fmt"
	"io"

	"github.com/golang/glog"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"

	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"math/big"
	"net/http"
	"time"

	"github.com/gorilla/mux"
)

const AI_DBNAME = "ai"

var AI_DB_CLIENT *gorm.DB

var (
	dbNameList = []string{
		"risk",
		"assets",
		"threaten",
		"control",
		AI_DBNAME,
	}
	dbCache = map[string]*gorm.DB{}
)

func GetRiskDB() *gorm.DB {
	return dbCache["risk"]
}

func GetAssetsDB() *gorm.DB {
	return dbCache["assets"]
}

func GetThreatenDB() *gorm.DB {
	return dbCache["threaten"]
}

func GetControlDB() *gorm.DB {
	return dbCache["control"]
}

func GetAiDB() *gorm.DB {
	return dbCache[AI_DBNAME]
}

func GetDB(name string) *gorm.DB {
	if db, ok := dbCache[name]; ok {
		return db
	}
	return nil
}

func InitPGDB() error {
	host := "10.254.177.177"
	// host := "10.254.179.93"
	port := 5432
	user := "postgres"
	password := "9).iz3ZS9)"
	// password := "9@iz3ZS9@"

	dsn := fmt.Sprintf("host=%s port=%d user=%s password=%s sslmode=disable", host, port, user, password)

	// 使用 pgx 作为底层驱动
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
	if err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}

	db.Exec(fmt.Sprintf("CREATE DATABASE %s;", AI_DBNAME))
	// if err != nil {
	//         glog.Error(err)
	//         return err
	// }

	for _, dbname := range dbNameList {
		dsn := fmt.Sprintf(
			"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
			host, port, user, password, dbname,
		)

		// 使用 pgx 作为底层驱动
		db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
		if err != nil {
			return fmt.Errorf("failed to connect to database: %w", err)
		}

		glog.Infof("✅ Successfully connected to PostgreSQL with GORM for %s!\n", dbname)
		dbCache[dbname] = db
	}
	return nil
}

type LLMCfgSt struct {
	ID          uint    `json:"id"` // 自增ID
	API         string  `json:"api"`
	Model       string  `json:"model"`
	User        string  `json:"user"`
	Token       string  `json:"token"`
	Temperature float32 `json:"temperature"`
	Think       bool    `json:"think"`
	Enable      bool    `json:"enable"`
}

func CreateLLMTable() error {
	err := GetAiDB().AutoMigrate(&LLMCfgSt{})
	if err != nil {
		glog.Error(err)
		return err
	}
	return nil
}

var PORT = "30445"

func generateSelfSignedCert() (tls.Certificate, error) {
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return tls.Certificate{}, err
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"Test Company"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(24 * time.Hour), // 有效期 24 小时
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return tls.Certificate{}, err
	}

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	return tls.X509KeyPair(certPEM, keyPEM)
}

func InitRestFullServer() {
	router := mux.NewRouter()

	// 大模型配置接口
	router.HandleFunc("/v1/llm/config", LLMCfgAddAPI).Methods("POST")       // 大模型配置-添加
	router.HandleFunc("/v1/llm/config", LLMCfgDelAPI).Methods("DELETE")     // 大模型配置-删除
	router.HandleFunc("/v1/llm/config", LLMCfgUpdateAPI).Methods("PUT")     // 大模型配置-修改
	router.HandleFunc("/v1/llm/config/list", LLMCfgListAPI).Methods("POST") // 大模型配置-列表

	cert, err := generateSelfSignedCert()
	if err != nil {
		panic(err)
	}

	server := &http.Server{
		Addr:    ":" + PORT,
		Handler: router,
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{cert},
		},
	}

	go server.ListenAndServeTLS("", "") // 证书已在 TLSConfig 中提供
}

type LLMCfgAPISt struct {
	ID          uint    `json:"id"`
	API         string  `json:"api"`
	Model       string  `json:"model"`
	User        string  `json:"user"`
	Token       string  `json:"token"`
	Temperature float32 `json:"temperature"`
	Think       bool    `json:"think"`
	Enable      bool    `json:"enable"`
}

type ApiResSt struct {
	Code    int         `json:"code"`
	Message string      `json:"message,omitempty"`
	Data    interface{} `json:"data,omitempty"`
}

func LLMCfgAddAPI(w http.ResponseWriter, r *http.Request) {
	req := LLMCfgAPISt{}
	// 读取请求体
	defer r.Body.Close()
	if r.Body == nil {
		glog.Error("request body is nil")
		return
	}
	body, _ := io.ReadAll(r.Body)
	err := json.Unmarshal(body, &req)

	glog.Info("body", &req)

	db := GetDB(AI_DBNAME)

	// cfg := LLMCfgSt(req)
	cfg := LLMCfgSt{
		API:         req.API,
		Model:       req.Model,
		User:        req.User,
		Token:       req.Token,
		Temperature: req.Temperature,
		Think:       req.Think,
		Enable:      req.Enable,
	}
	db.Create(&cfg)

	res := ApiResSt{Code: 0}
	resJson, _ := json.Marshal(res)
	w.Write(resJson)

	if err != nil {
		glog.Error(err)
		return
	}
}
func LLMCfgUpdateAPI(w http.ResponseWriter, r *http.Request) {
	req := LLMCfgAPISt{}
	// 读取请求体
	defer r.Body.Close()
	if r.Body == nil {
		glog.Error("request body is nil")
		return
	}
	body, _ := io.ReadAll(r.Body)
	err := json.Unmarshal(body, &req)

	glog.Info("body", req)

	db := GetDB(AI_DBNAME)

	// cfg := LLMCfgSt(req)
	db.Save(&LLMCfgSt{
		ID:          req.ID,
		API:         req.API,
		Model:       req.Model,
		User:        req.User,
		Token:       req.Token,
		Temperature: req.Temperature,
		Think:       req.Think,
		Enable:      req.Enable,
	})

	res := ApiResSt{Code: 0}
	resJson, _ := json.Marshal(res)
	w.Write(resJson)

	if err != nil {
		glog.Error(err)
		return
	}
}
func LLMCfgDelAPI(w http.ResponseWriter, r *http.Request) {
	req := LLMCfgAPISt{}
	// 读取请求体
	defer r.Body.Close()
	if r.Body == nil {
		glog.Error("request body is nil")
		return
	}
	body, _ := io.ReadAll(r.Body)
	err := json.Unmarshal(body, &req)

	db := GetDB(AI_DBNAME)
	db.Delete(&LLMCfgSt{}, req.ID)

	glog.Info("id", req.ID)

	if err != nil {
		glog.Error(err)
		return
	}
}
func LLMCfgListAPI(w http.ResponseWriter, r *http.Request) {
	req := LLMCfgAPISt{}
	// 读取请求体
	defer r.Body.Close()
	if r.Body == nil {
		glog.Error("request body is nil")
		return
	}
	body, _ := io.ReadAll(r.Body)
	err := json.Unmarshal(body, &req)

	db := GetDB(AI_DBNAME)

	var configs []LLMCfgSt
	db.Find(&configs)

	glog.Info("configs", configs)

	res := ApiResSt{Code: 0, Data: configs}
	resJson, _ := json.Marshal(res)
	w.Write(resJson)

	if err != nil {
		glog.Error(err)
		return
	}
}

func main() {
	flag.Parse()

	glog.Info()
	InitPGDB()
	CreateLLMTable()
	InitRestFullServer()
	select {}
}
